"""
可视化操作关联矩阵图 - 使用 Pyvis 构建有向图
顶点是字段，边是矩阵中的权重值，按照字段对应的 Operation 进行聚类
"""

import json
import argparse
from pathlib import Path
from pyvis.network import Network
from typing import Dict, List, Tuple, Set
from collections import defaultdict


class DependencyGraphVisualizer:
    """操作关联图可视化器"""

    def __init__(self, data_file: str, min_weight: float = 0.01):
        """
        初始化可视化器
        
        Args:
            data_file: 预处理结果JSON文件路径
            min_weight: 最小权重阈值，低于此值的边将被过滤
        """
        self.data_file = data_file
        self.min_weight = min_weight
        self.correlation_matrices = {}
        self.operations = set()
        self.fields_by_operation = defaultdict(set)

    def load_data(self):
        """加载预处理结果数据"""
        print(f"正在加载数据文件: {self.data_file}")

        with open(self.data_file, 'r', encoding='utf-8') as f:
            data = json.load(f)

        self.correlation_matrices = data.get('operation_correlation_matrices', {})
        total_matrices = sum(len(matrices) for matrices in self.correlation_matrices.values()) if isinstance(
            next(iter(self.correlation_matrices.values()), None), dict) else len(self.correlation_matrices)
        print(f"加载了 {total_matrices} 个操作关联矩阵")

        # 提取所有字段和操作
        self._extract_fields_and_operations()

    def _extract_fields_and_operations(self):
        """提取所有字段和对应的操作"""
        for process_or_key, process_matrices_or_data in self.correlation_matrices.items():
            # 格式为 {process: {matrix_key: matrix_data}}
            for matrix_key, matrix_data in process_matrices_or_data.items():
                self._process_matrix_data(matrix_data)

        print(f"发现 {len(self.operations)} 个操作")
        print(f"操作列表: {sorted(self.operations)}")

    def _process_matrix_data(self, matrix_data):
        """处理单个矩阵数据"""
        src_fields = matrix_data.get('src_fields', [])
        tgt_fields = matrix_data.get('tgt_fields', [])

        # 处理源字段
        for field in src_fields:
            operation = self._extract_operation_from_field(field)
            self.operations.add(operation)
            self.fields_by_operation[operation].add(field)

        # 处理目标字段
        for field in tgt_fields:
            operation = self._extract_operation_from_field(field)
            self.operations.add(operation)
            self.fields_by_operation[operation].add(field)

    def _extract_operation_from_field(self, field: str) -> str:
        """
        从字段名中提取操作名
        例如: checkout.placeOrder.address.city -> checkout.placeOrder
        """
        parts = field.split('.')
        if len(parts) >= 2:
            return f"{parts[0]}.{parts[1]}"
        return field

    def _get_operation_color(self, operation: str) -> str:
        """为不同操作分配颜色"""
        colors = [
            '#FF6B6B',  # 红色
            '#4ECDC4',  # 青色
            '#45B7D1',  # 蓝色
            '#96CEB4',  # 绿色
            '#FFEAA7',  # 黄色
            '#DDA0DD',  # 紫色
            '#98D8C8',  # 薄荷绿
            '#F7DC6F',  # 金色
            '#BB8FCE',  # 淡紫色
            '#85C1E9',  # 天蓝色
        ]

        # 根据操作名的哈希值选择颜色
        color_index = hash(operation) % len(colors)
        return colors[color_index]

    def _get_field_color(self, operation: str) -> str:
        """为字段节点分配颜色（比操作节点稍浅）"""
        base_color = self._get_operation_color(operation)
        # 将颜色转换为更浅的版本
        lighter_colors = {
            '#FF6B6B': '#FFB3B3',  # 浅红色
            '#4ECDC4': '#A3E6E1',  # 浅青色
            '#45B7D1': '#A2D5E8',  # 浅蓝色
            '#96CEB4': '#C8E6D9',  # 浅绿色
            '#FFEAA7': '#FFF4D3',  # 浅黄色
            '#DDA0DD': '#EECCEE',  # 浅紫色
            '#98D8C8': '#C4E8DD',  # 浅薄荷绿
            '#F7DC6F': '#FBECB7',  # 浅金色
            '#BB8FCE': '#DDC7E7',  # 浅淡紫色
            '#85C1E9': '#C2E0F4',  # 浅天蓝色
        }
        return lighter_colors.get(base_color, '#E0E0E0')

    def _calculate_operation_positions(self):
        """计算每个操作的基础位置，形成圆形分布"""
        positions = {}
        operations = list(self.operations)
        num_operations = len(operations)

        # 在圆形上分布操作
        radius = 400
        for i, operation in enumerate(operations):
            angle = (2 * 3.14159 * i) / num_operations
            x = radius * (1 + 0.5 * i) * (1 if i % 2 == 0 else -1)
            y = radius * (1 + 0.3 * i) * (1 if (i // 2) % 2 == 0 else -1)
            positions[operation] = (x, y)

        return positions

    def create_graph(self, output_file):
        """创建并保存图形"""
        print("正在创建图形...")

        # 创建网络图
        net = Network(
            height="800px",
            width="100%",
            bgcolor="#222222",
            font_color="white",
            directed=True
        )

        # 设置物理引擎参数
        net.set_options("""
        var options = {
          "physics": {
            "enabled": true,
            "stabilization": {"iterations": 100},
            "barnesHut": {
              "gravitationalConstant": -2000,
              "centralGravity": 0.3,
              "springLength": 95,
              "springConstant": 0.04,
              "damping": 0.09
            }
          },
          "nodes": {
            "font": {"size": 12}
          },
          "edges": {
            "font": {"size": 10},
            "smooth": {"type": "continuous"}
          }
        }
        """)

        # 添加节点（字段）
        self._add_nodes(net)

        # 添加边（关联关系）
        self._add_edges(net)

        # 保存图形
        net.save_graph(output_file)
        print(f"图形已保存到: {output_file}")

    def _add_nodes(self, net: Network):
        """添加节点到网络图 - 包括 Operation 节点和 Field 节点"""
        print("添加节点...")

        node_count = 0
        operation_positions = self._calculate_operation_positions()

        # 1. 添加 Operation 节点
        for operation in self.operations:
            base_x, base_y = operation_positions[operation]
            operation_color = self._get_operation_color(operation)

            # Operation 节点使用更大的尺寸和特殊形状
            net.add_node(
                f"op_{operation}",  # Operation 节点ID前缀
                label=operation.split('.')[-1],  # 只显示操作名，不包括服务名
                color={
                    'background': operation_color,
                    'border': '#ffffff',
                    'highlight': {'background': operation_color, 'border': '#ffff00'}
                },
                title=f"Operation: {operation}\nFields: {len(self.fields_by_operation[operation])}",
                size=40,  # Operation 节点更大
                shape='box',  # 使用方形区分
                font={'size': 16, 'color': 'white', 'bold': True},
                x=base_x,
                y=base_y,
                physics=True,
                mass=5  # Operation 节点质量更大
            )
            node_count += 1

        # 2. 添加 Field 节点
        for operation, fields in self.fields_by_operation.items():
            base_x, base_y = operation_positions[operation]
            field_color = self._get_field_color(operation)

            field_list = list(fields)
            for i, field in enumerate(field_list):
                # 简化字段名显示
                display_name = self._simplify_field_name(field)

                # 在操作节点周围圆形分布字段节点
                angle = (2 * 3.14159 * i) / len(field_list)
                radius = 120  # 字段节点距离操作节点的半径
                x = base_x + radius * (1.5 + 0.3 * len(field_list)) * (1 if i % 2 == 0 else -1)
                y = base_y + radius * (1 + 0.2 * i) * (1 if (i // 2) % 2 == 0 else -1)

                net.add_node(
                    field,  # 使用完整字段名作为ID
                    label=display_name,  # 显示简化名称
                    color={
                        'background': field_color,
                        'border': self._get_operation_color(operation),
                        'highlight': {'background': field_color, 'border': '#ffff00'}
                    },
                    title=f"Operation: {operation}\nField: {field}",
                    size=20,  # Field 节点较小
                    shape='dot',  # 使用圆形
                    font={'size': 10, 'color': 'white'},
                    x=x,
                    y=y,
                    physics=True,
                    mass=1
                )
                node_count += 1

        print(
            f"添加了 {node_count} 个节点 ({len(self.operations)} 个操作节点 + {node_count - len(self.operations)} 个字段节点)")

    def _simplify_field_name(self, field: str) -> str:
        """简化字段名显示"""
        parts = field.split('.')
        if len(parts) > 2:
            # 只显示最后两部分
            return '.'.join(parts[-2:])
        return field

    def _add_edges(self, net: Network):
        """添加边到网络图 - 包括从属关系虚线和字段关联实线"""
        print("添加边...")

        edge_count = 0

        # 1. 添加 Operation 到 Field 的从属关系虚线
        print("添加从属关系虚线...")
        ownership_count = 0
        for operation, fields in self.fields_by_operation.items():
            operation_node_id = f"op_{operation}"

            for field in fields:
                net.add_edge(
                    operation_node_id,
                    field,
                    color={'color': '#888888', 'opacity': 0.6},  # 灰色虚线
                    width=1,
                    dashes=[5, 5],  # 虚线样式
                    title=f"Ownership: {operation} owns {field}",
                    physics=True,
                    smooth={'type': 'continuous'}
                )
                ownership_count += 1

        print(f"添加了 {ownership_count} 条从属关系虚线")

        # 2. 添加 Field 到 Field 的关联关系实线
        print("添加字段关联实线...")
        correlation_count = 0
        for process_or_key, process_matrices_or_data in self.correlation_matrices.items():
            # 格式为 {process: {matrix_key: matrix_data}}
            for matrix_key, matrix_data in process_matrices_or_data.items():
                correlation_count += self._add_field_correlations(net, matrix_key, matrix_data)

        print(f"添加了 {correlation_count} 条字段关联实线")

    def _add_field_correlations(self, net, matrix_key, matrix_data):
        """添加字段关联关系，返回添加的关联数"""
        count = 0
        src_fields = matrix_data.get('src_fields', [])
        tgt_fields = matrix_data.get('tgt_fields', [])
        correlation_matrix = matrix_data.get('correlation_matrix', [])

        # 检查矩阵维度
        if len(correlation_matrix) != len(src_fields):
            print(f"警告: 矩阵行数 ({len(correlation_matrix)}) 与源字段数 ({len(src_fields)}) 不匹配")
            return count

        # 遍历关联矩阵
        for i, src_field in enumerate(src_fields):
            if i >= len(correlation_matrix):
                continue

            row = correlation_matrix[i]
            if len(row) != len(tgt_fields):
                print(f"警告: 矩阵列数 ({len(row)}) 与目标字段数 ({len(tgt_fields)}) 不匹配")
                continue

            for j, tgt_field in enumerate(tgt_fields):
                if j >= len(row):
                    continue

                weight = row[j]

                # 过滤低权重边
                if weight >= self.min_weight:
                    # 计算边的宽度和颜色
                    width = max(2, weight * 15)  # 权重越大，边越粗

                    # 根据权重设置边的颜色
                    if weight > 0.1:
                        edge_color = "#FF4444"  # 高权重：红色
                    elif weight > 0.05:
                        edge_color = "#FFAA44"  # 中权重：橙色
                    else:
                        edge_color = "#44AAFF"  # 低权重：蓝色

                    net.add_edge(
                        src_field,
                        tgt_field,
                        width=width,
                        color={'color': edge_color, 'opacity': 0.8},
                        title=f"Correlation: {weight:.4f}\n{matrix_key}",
                        label=f"{weight:.3f}",
                        arrows={'to': {'enabled': True, 'scaleFactor': 1.2}},
                        physics=True,
                        smooth={'type': 'continuous'}
                    )
                    count += 1

        return count

    def generate_statistics(self):
        """生成统计信息"""
        print("\n" + "=" * 60)
        print("统计信息")
        print("=" * 60)

        total_fields = sum(len(fields) for fields in self.fields_by_operation.values())
        print(f"总操作数: {len(self.operations)}")
        print(f"总字段数: {total_fields}")

        total_matrices = 0
        for process_or_key, process_matrices_or_data in self.correlation_matrices.items():
            total_matrices += len(process_matrices_or_data)
        print(f"总关联矩阵数: {total_matrices}")

        print(f"\n各操作的字段数:")
        for operation in sorted(self.operations):
            field_count = len(self.fields_by_operation[operation])
            print(f"  {operation}: {field_count} 个字段")

        # 统计边的权重分布
        all_weights = []
        for process_or_key, process_matrices_or_data in self.correlation_matrices.items():
            # 格式为 {process: {matrix_key: matrix_data}}
            for matrix_key, matrix_data in process_matrices_or_data.items():
                correlation_matrix = matrix_data.get('correlation_matrix', [])
                for row in correlation_matrix:
                    for weight in row:
                        if weight > 0:
                            all_weights.append(weight)

        if all_weights:
            print(f"\n权重统计:")
            print(f"  非零权重数: {len(all_weights)}")
            print(f"  最大权重: {max(all_weights):.4f}")
            print(f"  最小权重: {min(all_weights):.4f}")
            print(f"  平均权重: {sum(all_weights) / len(all_weights):.4f}")


if __name__ == '__main__':
    """主函数"""
    parser = argparse.ArgumentParser(description='可视化操作关联矩阵图')
    parser.add_argument('--data_file', type=str,
                        default='data/cbt/checkout_100_user50_preprocess_result.json',
                        help='预处理结果JSON文件路径')
    parser.add_argument('--output', type=str,
                        default='data/dependency_graph.html',
                        help='输出HTML文件名')
    parser.add_argument('--min_weight', type=float,
                        default=0.0001,
                        help='最小权重阈值')

    args = parser.parse_args()

    # 检查输入文件
    if not Path(args.data_file).exists():
        print(f"错误: 数据文件 {args.data_file} 不存在")
        exit(1)

    # 创建可视化器
    visualizer = DependencyGraphVisualizer(args.data_file, args.min_weight)

    try:
        # 加载数据
        visualizer.load_data()

        # 生成统计信息
        visualizer.generate_statistics()

        # 创建图形
        visualizer.create_graph(args.output)

        print(f"\n可视化完成! 请在浏览器中打开 {args.output} 查看结果")

    except Exception as e:
        print(f"执行过程中出现错误: {e}")
        import traceback

        traceback.print_exc()
